import pandas as pd
import io
import dataclasses
import re
import os
from collections import defaultdict
from colorama import init, Fore
from draw_plot import draw_plot

def info(str):
    print(Fore.GREEN + str + Fore.RESET)
init()  # Initialize colorama
@dataclasses.dataclass
class CSV:
    data:pd.DataFrame
    hz:int
    temprature:float
    def to_csv(self,file_path:str,encoding:str='GBK'):
        self.data.to_csv(file_path, index=False,encoding=encoding)

def extract_data_from_file(file_path)->CSV:
    
    data_str = ""
    start_collecting = False
    
    # Extract temperature from filename patterns
    temp_pattern1 = r'(\d+)-\d+V-\d+Hz.csv$'  # For pattern like "30-30V-400Hz.csv"
    temp_pattern2 = r'(\d+)T.csv$'            # For pattern like "30T.csv"

    match1 = re.search(temp_pattern1, file_path)
    match2 = re.search(temp_pattern2, file_path)

    if match1:
        temprature = float(match1.group(1))
    elif match2:
        temprature = float(match2.group(1))
    else:
        assert False, f"Invalid filename pattern: {file_path}"
    
    hz = 0
    
    with open(file_path, 'r') as file:
        for line in file:
            # Strip whitespace from the line
            line = line.strip()
            
            if "Hz" in line and "Frequency" in line:
                idx = line.find("Frequency")
                end = line.find("Hz")
                hz = int(line[idx+10:end])
            
            # Check if line contains 'aaaaaa' to start collecting
            if 'Elapsed Time[sec]' in line:
                start_collecting = True
                data_str += line + '\n'
                continue
            
            # If we've started collecting and encounter 'cycle', stop
            elif start_collecting and 'Cycle' in line:
                break
            
            # If we're collecting and line is not empty, add it to data_str
            elif start_collecting and line:
                data_str += line + '\n'

    # Convert collected string to DataFrame
    if data_str:
        # Using StringIO to convert string to DataFrame
        df = pd.read_csv(io.StringIO(data_str), na_values=['-'])
        csv = CSV(data=df,hz=hz,temprature=temprature)
        return csv
    else:
        raise Exception("No data found in the file")

def retain_head(csv: CSV, head_list: list[str],remove_nan_line:bool=True) -> CSV:
    # 保留指定的列
    retained_data = csv.data[head_list]
    # 如果remove_nan_line为True，删除包含空值的行
    if remove_nan_line:
        retained_data = retained_data.dropna()
    # 返回新的CSV对象
    return CSV(data=retained_data, hz=csv.hz, temprature=csv.temprature)
    
def to_float(csv: CSV) -> CSV:
    # 将数据部分转换为浮点数
    data = csv.data.astype(float)
    # 返回新的CSV对象
    return CSV(data=data, hz=csv.hz, temprature=csv.temprature)

def sort_by_column(csv: CSV, column_name: str="Elapsed Time[sec]") -> CSV:
    # 按照指定列排序
    sorted_data = csv.data.sort_values(by=column_name)
    # 返回新的CSV对象
    return CSV(data=sorted_data, hz=csv.hz, temprature=csv.temprature)

def get_dir_files(dir_path:str,regex:str='(\d+)T.csv$')->list[str]:
    files = os.listdir(dir_path)
    files = [os.path.join(dir_path, file) for file in files if re.match(regex, file)]
    return files

def merge_by_hz(dir_path:str,out_put_dir:str=".",list_head:list[str]=['Elapsed Time[sec]', 'Current[A]']):
    files = get_dir_files(dir_path)
    info(f"共找到{len(files)}个文件,开始合并{list_head}列")
    data = defaultdict(list)
    for file in files:
        csv = extract_data_from_file(file)
        csv = retain_head(csv, list_head)
        csv = to_float(csv)
        csv = sort_by_column(csv)
        new_col = {}
        for col in csv.data.columns:
            new_col[col] = f"{col} {csv.temprature}T"
        csv.data.rename(columns=new_col,inplace=True)
        data[csv.hz].append(csv.data)
    info("数据提取完成，开始合并")
    for hz in data:
        a = data[hz]
        frame = pd.concat(data[hz],axis=1)
        file_path = os.path.join(out_put_dir,f"{hz}Hz.csv")
        frame.to_csv(f"{file_path}",index=False,encoding="GBK")
        info(f"合并{hz}Hz数据到{file_path},共{len(data[hz])}个文件")
    info("合并完成")

def get_T(file_path:str)->list[str]:
    csv = pd.read_csv(file_path,encoding="GBK")
    list_head = csv.columns
    m = {}
    rex = r'(\d+)T$'
    for head in list_head:
        head = head.split(" ")
        for h in head:
            match = re.search(rex, h)
            if match and "T" in h:
                m[h] = 1
    return m.keys()


def draw_hz_csv(file_path:str,need_T:list[str],list_head:list[str]=['Elapsed Time[sec]', 'Current[A]'],point_size:float=3):
    rex = r'(\d+)Hz.csv$'
    match = re.search(rex, file_path)
    if not match:
        raise Exception(f"Invalid filename pattern: {file_path}")
    hz = int(match.group(1))
    csv = pd.read_csv(file_path,encoding="GBK")
    x = []
    y = []
    group = []
    for col in csv.columns:
        precol = col
        if list_head[0] in col:
            col = col.replace(list_head[0],"").replace(" ","")
            if col not in need_T:
                continue
            x.append(csv[precol].values)
        else:
            col = col.replace(list_head[1],"").replace(" ","")
            if col not in need_T:
                continue
            y.append(csv[precol].values)
            group.append(col)
    draw_plot(x,y,group,label_x=list_head[0],label_y=list_head[1],title=f"{hz}Hz",point_size=point_size)

if __name__ == '__main__':
    
    rex = r"(\d+)T.csv$"
    print(re.match(rex, "30T.csv"))
    
    # Example usage
    file_path = 'C:/MyProject/cpp_project/for_lxj/tmp/C1Bn-PPQPCN2.0/35T.csv'  # Replace with your actual file path
    file_path = 'C:/MyProject/cpp_project/for_lxj/data/C4-DIO-CF2O-Ph-Ph-CF2O-Ph-CN/30-30V-400Hz.csv'
    try:
        a = [[2,1,3]]
        b = [[4,5,6]]
        # draw_plot(a,b,["a","b"])
        # merge_by_hz("C:/MyProject/cpp_project/for_lxj/tmp/C1Bn-PPQPCN2.0")
        rex = r'(\d+)T$'
        # print(re.search(rex, "30.0T"))
        print(get_T('1000Hz.csv'))
        
    except Exception as e:
        print(f"Error processing file: {str(e)}")
